# +----------------------------------------------------------------------
# | ChatWork智能聊天办公系统
# +----------------------------------------------------------------------
# | 软件声明: 本系统并非自由软件,未经授权任何形式的商业使用均属非法。
# | 版权保护: 任何企业和个人不允许对程序代码以任何形式任何目的复制/分发。
# | 授权要求: 如有商业使用需求,请务必先与版权所有者取得联系并获得正式授权。
# +----------------------------------------------------------------------
# | Author: ChatWork Team <2474369941@qq.com>
# +----------------------------------------------------------------------
from tortoise import fields
from kernels.model import DbModel


class AgentAttachmentModel(DbModel):
    id = fields.CharField(pk=True, max_length=32, description="主键")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    file_md5 = fields.CharField(null=False, max_length=32, default="", description="文件md5")
    file_type = fields.CharField(null=False, max_length=8, default="", description="文件类型")
    file_name = fields.CharField(null=False, max_length=200, default="", description="文件名称")
    file_size = fields.IntField(null=False, default=0, description="文件大小")
    file_ext = fields.CharField(null=False, max_length=10, default="", description="文件大小")
    file_path = fields.CharField(null=False, max_length=250, default="", description="文件路径")
    text_path = fields.CharField(null=False, max_length=250, default="", description="文本路径")
    ip = fields.CharField(null=False, max_length=64, default="", description="来源IP")
    create_time = fields.IntField(null=False, default=0, description="创建时间")

    class Meta:
        table_description = "会话附件表"
        table = DbModel.table_prefix("agent_attachment")


class AgentChatRecordsModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    scene = fields.CharField(null=False, max_length=10, default="normal", description="场景: [normal,test,web,api]")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    robot_id = fields.IntField(null=False, default=0, description="智脑ID")
    session_id = fields.IntField(null=False, default=0, description="会话ID")
    publish_id = fields.IntField(null=False, default=0, description="发布ID")
    publish_uid = fields.CharField(null=False, max_length=32, default="", description="发布用户")
    use_tokens = fields.IntField(null=False, default=0, description="消耗Tokens")
    use_points = fields.IntField(null=False, default=0, description="消耗的积分")
    question = fields.TextField(description="提问问题")
    answer = fields.TextField(description="答复内容")
    thinks = fields.TextField(description="思考内容")
    context = fields.TextField(description="上下文值")
    quotes = fields.TextField(description="引用内容")
    attach = fields.TextField(description="答复附件")
    annex = fields.TextField(description="提问附件")
    issues = fields.CharField(null=False, max_length=800, default="", description="相关问题")
    tools_list = fields.TextField(description="MCP工具")
    web_search = fields.TextField(description="联网搜索")
    flows_usage = fields.TextField(description="消耗流程")
    ip = fields.CharField(null=False, max_length=32, default="", description="来源IP")
    end_clear = fields.SmallIntField(null=False, default=0, description="清除末端: [0=否, 1=是]")
    is_clear = fields.SmallIntField(null=False, default=0, description="清除记录: [0=否, 1=是]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    user_delete = fields.SmallIntField(null=False, default=0, description="用户删除: [0=否, 1=是]")
    task_time = fields.CharField(null=False, max_length=20, default=0, description="执行耗时")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "对话记录表"
        table = DbModel.table_prefix("agent_chat_records")


class AgentChatRoomingModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    bot = fields.CharField(null=False, max_length=32, default="", description="智能体")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    robot_id = fields.IntField(null=False, default=0, description="智能体ID")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "对话应用表"
        table = DbModel.table_prefix("agent_chat_rooming")


class AgentChatSessionModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    room_id = fields.IntField(null=False, default=0, description="房间ID")
    robot_id = fields.IntField(null=False, default=0, description="智能体ID")
    name = fields.TextField(null=False, default="", description="会话名称")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "对话会话表"
        table = DbModel.table_prefix("agent_chat_session")


class AgentRobotModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    type = fields.CharField(null=False, max_length=32, default="", description="类型: [simple=简易应用, workflow=工作流]")
    code = fields.CharField(null=False, max_length=32, default="", description="编号")
    user_id = fields.IntField(null=False, default=0, description="所属用户")
    mcp_ids = fields.CharField(null=False, max_length=250, default="", description="关联MCP")
    know_ids = fields.CharField(null=False, max_length=250, default="", description="关联知识库")
    logo = fields.CharField(null=False, max_length=250, default="", description="智能体图标")
    name = fields.CharField(null=False, max_length=100, default="", description="智能体名称")
    intro = fields.CharField(null=False, max_length=800, default="", description="智能体简介")
    llm_models = fields.CharField(null=False, max_length=100, default="", description="对话的模型")
    llm_config = fields.CharField(null=False, max_length=900, default="", description="模型的参数")
    search_mode = fields.CharField(null=False, max_length=20, default="", description="检索的模式: [similar,full,mix]")
    search_tokens = fields.IntField(null=False, default=5000, description="检索引用上限: [100-30000]")
    search_similar = fields.DecimalField(null=False, max_digits=6, decimal_places=3, default=0, description="检索相似度")
    annex_similar = fields.DecimalField(null=False, max_digits=6, decimal_places=3, default=0, description="附件相似度")
    ranking_status = fields.SmallIntField(null=False, default=0, description="检索的重排: [0=否, 1=是]")
    ranking_score = fields.DecimalField(null=False, max_digits=6, decimal_places=3, default=0, description="重排分数")
    ranking_model = fields.CharField(null=False, max_length=100, default="", description="重排模型")
    optimize_ask = fields.SmallIntField(null=False, default=0, description="问题优化: [0=否, 1=是]")
    optimize_model = fields.CharField(null=False, max_length=100, default="", description="优化模型")
    prompt_roles = fields.TextField(default="", description="角色提示词")
    prompt_limit = fields.TextField(default="", description="限定提示词")
    welcome_text = fields.TextField(default="", description="欢迎光临词")
    sort = fields.IntField(null=False, default=0, description="排序编号")
    broadcast = fields.CharField(null=False, default="", max_length=100, description="语音播报参")
    question_guide = fields.SmallIntField(null=False, default=0, description="猜您想问的")
    empty_search_text = fields.TextField(null=False, default="", description="空搜索文本")
    empty_search_type = fields.CharField(null=False, max_length=10, default="ai", description="空搜索类型: [ai,default]")
    max_upload_num = fields.SmallIntField(null=False, default=0, description="最大上传数: [0-15]")
    is_upload_files = fields.SmallIntField(null=False, default=0, description="文档上传: [0=否, 1=是]")
    is_upload_image = fields.SmallIntField(null=False, default=0, description="图片上传: [0=否, 1=是]")
    is_show_timed = fields.SmallIntField(null=False, default=0, description="显示耗时: [0=否, 1=是]")
    is_show_quote = fields.SmallIntField(null=False, default=0, description="显示引用: [0=否, 1=是]")
    is_show_context = fields.SmallIntField(null=False, default=0, description="显上下文: [0=否, 1=是]")
    is_original_reply = fields.SmallIntField(null=False, default=0, description="原文答复: [0=否, 1=是]")
    is_disable = fields.SmallIntField(null=False, default=0, description="是否禁用: [0=否, 1=是]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "智能体表"
        table = DbModel.table_prefix("agent_robot")

    @classmethod
    async def queryRobotDictsById(cls, ids: list, field: list = None) -> dict:
        """ 获取智能体字典 """
        dicts = {}
        if ids:
            field = ["id", "code", "logo", "name", "llm_models"] if not field else field
            robots = await (cls.filter(id__in=list(set(ids))).all().values(*field))
            dicts = {item["id"]: item for item in robots}
        return dicts


class AgentRobotCollectModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    robot_id = fields.IntField(null=False, default=0, description="智能体ID")
    total_num = fields.IntField(null=False, default=0, description="合计次数")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "智脑收藏表"
        table = DbModel.table_prefix("agent_robot_collect")


class AgentRobotInstructModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    robot_id = fields.IntField(null=False, default=0, description="智能体ID")
    scene = fields.CharField(null=False, max_length=10, default="", description="场景: [fixed=固定, direct=指令]")
    keyword = fields.CharField(null=False, max_length=200, default="", description="关键词")
    content = fields.TextField(default="", description="回复内容")
    annex = fields.TextField(default="", description="上传附件")
    sort = fields.IntField(null=False, default=0, description="排序编号")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")

    class Meta:
        table_description = "智能体指令表"
        table = DbModel.table_prefix("agent_robot_instruct")


class AgentRobotPublishModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    scene = fields.CharField(null=False, max_length=20, default="", description="场景值: [web, api]")
    robot_id = fields.IntField(null=False, default=0, description="智能体ID")
    api_key = fields.CharField(null=False, max_length=100, default="", description="密钥")
    api_secret = fields.CharField(null=False, max_length=100, default="", description="密码")
    name = fields.CharField(null=False, max_length=100, default="", description="名称")
    rpm_limit = fields.IntField(null=False, default=0, description="每分钟限制")
    max_points = fields.IntField(null=False, default=0, description="最大额度")
    use_points = fields.IntField(null=False, default=0, description="消耗额度")
    call_count = fields.IntField(null=False, default=0, description="调用次数")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    last_ip = fields.CharField(null=False, max_length=32, default="", description="最后使用的IP")
    last_time = fields.IntField(null=False, default=0, description="最后使用时间")
    expire_time = fields.IntField(null=False, default=0, description="创建时间")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "智能体发布表"
        table = DbModel.table_prefix("agent_robot_publish")


class AgentRobotTeamedModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    robot_id = fields.IntField(null=False, default=0, description="智能体ID")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    power = fields.SmallIntField(null=False, default=0, description="权力: [1=查看者, 2=编辑者, 3-管理者]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")

    class Meta:
        table_description = "智能体团队表"
        table = DbModel.table_prefix("agent_robot_teamed")


class AgentKnowModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    user_id = fields.IntField(null=False, default=0, description="所属用户")
    code = fields.CharField(null=False, max_length=32, default="", description="编号")
    logo = fields.CharField(null=False, max_length=250, default="", description="封面")
    name = fields.CharField(null=False, max_length=100, default="", description="名称")
    intro = fields.CharField(null=False, max_length=800, default="", description="简介")
    vector_model = fields.CharField(null=False, max_length=60, default="", description="向量模型")
    splits_model = fields.CharField(null=False, max_length=60, default="", description="拆分模型")
    sort = fields.IntField(null=False, default=0, description="排序")
    is_disable = fields.SmallIntField(null=False, default=0, description="是否禁用: [0=否, 1=是]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "知识库表"
        table = DbModel.table_prefix("agent_know")

    @classmethod
    async def queryKnowDictsById(cls, ids: list, field: list = None) -> dict:
        """ 获取知识库字典 """
        dicts = {}
        if ids:
            field = ["id", "code", "logo", "name", "vector_model", "splits_model"] if not field else field
            users = await (cls.filter(id__in=list(set(ids))).all().values(*field))
            dicts = {item["id"]: item for item in users}
        return dicts


class AgentKnowArchiveModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    know_id = fields.IntField(null=False, default=0, description="知识库ID")
    user_id = fields.IntField(null=False, default=0, description="拥有者UID")
    last_uid = fields.IntField(null=False, default=0, description="修改者UID")
    scene = fields.CharField(null=False, max_length=10, default="", description="来源场景: [file, text, link]")
    code = fields.CharField(null=False, max_length=32, default="", description="文档编号")
    name = fields.CharField(null=False, max_length=200, default="", description="文件名称")
    path = fields.CharField(null=False, max_length=250, default="", description="文件路径")
    size = fields.IntField(null=False, default=0, description="文件大小")
    train_mode = fields.CharField(null=False, max_length=10, default="", description="训练模式: [chunk=分段, qa=拆分]")
    separators = fields.CharField(null=False, max_length=200, default="", description="分段规则")
    chunk_size = fields.IntField(null=False, default=0, description="分段长度")
    chunk_overlap = fields.IntField(null=False, default=0, description="重叠长度")
    fail_reason = fields.TextField(default="", description="失败原因")
    qa_status = fields.SmallIntField(null=False, default=0, description="拆分状态: [0=等待, 1=排队, 2=处理, 3=成功, 4=失败]")
    is_disable = fields.SmallIntField(null=False, default=0, description="是否禁用: [0=否, 1=是]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "知识库文档表"
        table = DbModel.table_prefix("agent_know_archive")

    @classmethod
    async def queryArchiveDictsById(cls, ids: list, field: list = None) -> dict:
        """ 获取文档字典 """
        dicts = {}
        if ids:
            field = ["id", "code", "name", "path", "size"] if not field else field
            lists = await (cls.filter(id__in=list(set(ids))).all().values(*field))
            dicts = {item["id"]: item for item in lists}
        return dicts


class AgentKnowSplittingModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    know_id = fields.IntField(null=False, default=0, description="知识库ID")
    archive_id = fields.IntField(null=False, default=0, description="文档的ID")
    import_uid = fields.IntField(null=False, default=0, description="导入用户")
    model_alias = fields.CharField(null=False, default="", max_length=200, description="模型别名")
    model_name = fields.CharField(null=False, default="", max_length=200, description="名称名称")
    file_name = fields.CharField(null=False, default="", max_length=500, description="文件名称")
    file_path = fields.CharField(null=False, default="", max_length=250, description="文件路径")
    file_size = fields.IntField(null=False, default=0, description="文件大小")
    content = fields.TextField(default="", description="段落内容")
    results = fields.TextField(default="", description="拆分结果: [json]")
    fail_reason = fields.TextField(default="", description="拆分结果: [json]")
    usage_points = fields.IntField(null=False, default=0, description="消耗积分")
    usage_tokens = fields.CharField(default="", max_length=500, description="消耗tokens")
    retry = fields.SmallIntField(null=False, default=0, description="重试次数")
    status = fields.SmallIntField(null=False, default=0, description="拆分状态: [0=待拆分, 2=拆分中, 3=成功, 4=失败]")
    task_time = fields.IntField(null=False, default=0, description="执行时间(s)")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")

    class Meta:
        table_description = "知识拆分表"
        table = DbModel.table_prefix("agent_know_splitting")


class AgentKnowTeamedModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    know_id = fields.IntField(null=False, default=0, description="知识库ID")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    power = fields.SmallIntField(null=False, default=0, description="权力: [1=查看者, 2=编辑者, 3-管理者]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")

    class Meta:
        table_description = "知识库团队表"
        table = DbModel.table_prefix("agent_know_teamed")


class AgentMcpModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    code = fields.CharField(null=False, max_length=32, default="", description="编号")
    logo = fields.CharField(null=False, max_length=250, default="", description="封面")
    name = fields.CharField(null=False, max_length=100, default="", description="名称")
    intro = fields.CharField(null=False, max_length=800, default="", description="简介")
    url = fields.CharField(null=False, max_length=900, default="", description="MCP地址")
    tools = fields.TextField(default="", description="MCP工具")
    sort = fields.IntField(null=False, default=0, description="排序编号")
    is_disable = fields.SmallIntField(null=False, default=0, description="是否禁用: [0=否, 1=是]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "MCP服务表"
        table = DbModel.table_prefix("agent_mcp")


class AgentMcpTeamedModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    mcp_id = fields.IntField(null=False, default=0, description="MCP_ID")
    user_id = fields.IntField(null=False, default=0, description="用户ID")
    power = fields.SmallIntField(null=False, default=0, description="权力: [1=查看者, 2=编辑者, 3-管理者]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")

    class Meta:
        table_description = "MCP团队表"
        table = DbModel.table_prefix("agent_mcp_teamed")


class AgentSquareModel(DbModel):
    id = fields.IntField(pk=True, description="主键")
    cate_id = fields.IntField(null=False, default=0, description="分类ID")
    robot_id = fields.IntField(null=False, default=0, description="智能体ID")
    user_id = fields.IntField(null=False, default=0, description="公开用户")
    sort = fields.IntField(null=False, default=0, description="排序编号")
    audit_result = fields.CharField(null=False, max_length=300, default="", description="审核结果")
    audit_status = fields.SmallIntField(null=False, default=0, description="审核状态: [0=待审, 1=审核通过, 2=审核拒绝]")
    is_enable = fields.SmallIntField(null=False, default=0, description="是否显示: [0=否, 1=是]")
    is_delete = fields.SmallIntField(null=False, default=0, description="是否删除: [0=否, 1=是]")
    create_time = fields.IntField(null=False, default=0, description="创建时间")
    update_time = fields.IntField(null=False, default=0, description="更新时间")
    delete_time = fields.IntField(null=False, default=0, description="删除时间")

    class Meta:
        table_description = "智能体广场表"
        table = DbModel.table_prefix("agent_square")

    @classmethod
    async def querySquareDictsByRobotId(cls, ids: list, field: list = None) -> dict:
        """ 获取广场字典 """
        dicts = {}
        if ids:
            field = cls.without_field("is_delete,delete_time") if not field else field
            result = await (cls.filter(robot_id__in=list(set(ids))).filter(is_delete=0).all().values(*field))
            dicts = {item["robot_id"]: item for item in result}
        return dicts
